# -*- coding: utf-8 -*-
# @Time    : 2021/3/16 14:40
# @Author  : huangwei
# @File    : table.py
# @Software: PyCharm
from PIL import Image
from skimage import measure
from table_method import resize_img, exp, get_line, draw_lines, dist, line_line, minAreaRectBox, resize, sqrt
import numpy as np
import cv2
import os
import math


def table_predict( img, tableNet, re_size ):
    """ 将图片传入模型中得到输出 """
    imgResize, fx, fy, dx, dy = resize_img(img, re_size)

    imgResize = np.array(imgResize)
    img_w, img_h = re_size

    # 对图像进行减均值、缩放以及交换通道
    image = cv2.dnn.blobFromImage(imgResize, 1, size=(img_w, img_h), swapRB=False)

    image = np.array(image) / 255
    tableNet.setInput(image)
    out = tableNet.forward()
    out = exp(out[0])
    out = out[:, dy:, dx:]
    return out, fx, fy


def get_table_rowcols( img, model, input_size, prob=0.1, alph=15, row=30, col=30 ):
    resize_img, fx, fy = resize(img, input_size)

    pred = model.predict(np.array([np.array(resize_img) / 255.0]))
    pred = pred[0]

    rows = pred[..., 0]
    cols = pred[..., 1]

    labels = measure.label(rows > prob, connectivity=2)

    # 画出得到的连通域
    # labels1 = labels
    # for m in range(labels1.shape[0]):
    #     for n in range(labels1.shape[1]):
    #         if labels1[m][n] != 0:
    #             labels1[m][n] = 200
    # cv2.imwrite('temp_path/ee.png', labels1)

    # 返回所有连通区域的属性列表
    regions = measure.regionprops(labels)
    # 该连通区域的边界外接框宽度大于阈值row则将该外接框的横向中心线加入到集合中
    RowsLines = [get_line(line.coords) for line in regions if line.bbox[3] - line.bbox[1] > row]

    # 同理找出所有竖线
    labels = measure.label(cols > prob, connectivity=2)
    regions = measure.regionprops(labels)
    ColsLines = [get_line(line.coords) for line in regions if line.bbox[2] - line.bbox[0] > col]

    ccolbox = []
    crowlbox = []
    if len(RowsLines) > 0:
        RowsLines = np.array(RowsLines)
        RowsLines[:, [0, 2]] = RowsLines[:, [0, 2]] / fx
        RowsLines[:, [1, 3]] = RowsLines[:, [1, 3]] / fy
        xmin = RowsLines[:, [0, 2]].min()
        xmax = RowsLines[:, [0, 2]].max()
        ymin = RowsLines[:, [1, 3]].min()
        ymax = RowsLines[:, [1, 3]].max()
        ccolbox = [[xmin, ymin, xmin, ymax], [xmax, ymin, xmax, ymax]]
        RowsLines = RowsLines.tolist()

    if len(ColsLines) > 0:
        ColsLines = np.array(ColsLines)
        ColsLines[:, [0, 2]] = ColsLines[:, [0, 2]] / fx
        ColsLines[:, [1, 3]] = ColsLines[:, [1, 3]] / fy

        xmin = ColsLines[:, [0, 2]].min()
        xmax = ColsLines[:, [0, 2]].max()
        ymin = ColsLines[:, [1, 3]].min()
        ymax = ColsLines[:, [1, 3]].max()
        ColsLines = ColsLines.tolist()
        crowlbox = [[xmin, ymin, xmax, ymin], [xmin, ymax, xmax, ymax]]

    RowsLines += crowlbox
    ColsLines += ccolbox

    row_box, col_box = adjust_lines(RowsLines, ColsLines, alph=alph)
    RowsLines += row_box
    ColsLines += col_box

    for i in range(len(RowsLines)):
        for j in range(len(ColsLines)):
            RowsLines[i] = line_line(RowsLines[i], ColsLines[j], alph)
            ColsLines[j] = line_line(ColsLines[j], RowsLines[i], alph)

    # # 将所有找出的线在黑色图片上画出，找出只有这些线存在的图片上的连通区域
    # tmp = np.zeros(input_size[::-1], dtype='uint8')
    # tmp = draw_lines(tmp, ColsLines + RowsLines, color=255, line_width=1)
    # cv2.imwrite('temp_path/aa.png', tmp)
    #
    # labels = measure.label(tmp > 0, connectivity=2)
    # regions = measure.regionprops(labels)
    #
    # # 去除这些连通区域中有一边较小的区域以及不在整个大连通区域的线段
    # for region in regions:
    #     ymin, xmin, ymax, xmax = region.bbox
    #     label = region.label
    #     if ymax - ymin < 20 or xmax - xmin < 20:
    #         labels[labels == label] = 0
    # labels = measure.label(labels > 0, connectivity=2)
    #
    # indY, indX = np.where(labels > 0)
    # xmin, xmax = indX.min(), indX.max()
    # ymin, ymax = indY.min(), indY.max()
    # RowsLines = [p for p in RowsLines if
    #              xmin <= p[0] <= xmax and xmin <= p[2] <= xmax and ymin <= p[1] <= ymax and ymin <= p[3] <= ymax]
    # ColsLines = [p for p in ColsLines if
    #              xmin <= p[0] <= xmax and xmin <= p[2] <= xmax and ymin <= p[1] <= ymax and ymin <= p[3] <= ymax]
    #
    # tmp = np.zeros(input_size[::-1], dtype='uint8')
    # tmp = draw_lines(tmp, ColsLines + RowsLines, color=255, line_width=1)
    # cv2.imwrite('temp_path/bb.png', tmp)
    #
    # # 将坐标转换成原图中的坐标
    # RowsLines = [[box[0] / fx, box[1] / fy, box[2] / fx, box[3] / fy] for box in RowsLines]
    # ColsLines = [[box[0] / fx, box[1] / fy, box[2] / fx, box[3] / fy] for box in ColsLines]
    return RowsLines, ColsLines


# def get_table_rowcols( img, tableNet, input_size, prob, row=100, col=100 ):
#     out, fx, fy = table_predict(img, tableNet, input_size)
#
#     rows = out[0]
#     cols = out[1]
#
#     # 实现连通区域标记，connectivity=2 表示8连通，返回的值表示与周围八个连通的块数
#     labels = measure.label(rows > prob, connectivity=2)
#
#     # labels1 = labels
#     # for m in range(labels1.shape[0]):
#     #     for n in range(labels1.shape[1]):
#     #         if labels1[m][n] != 0:
#     #             labels1[m][n] = 200
#     # cv2.imwrite('temp_path/ee.png', labels1)
#
#     # 返回所有连通区域的属性列表
#     regions = measure.regionprops(labels)
#     # 该连通区域的边界外接框宽度大于阈值row则将该外接框的横向中心线加入到集合中
#     RowsLines = [get_line(line.coords) for line in regions if line.bbox[3] - line.bbox[1] > row]
#
#     # 同理找出所有竖线
#     labels = measure.label(cols > prob, connectivity=2)
#     regions = measure.regionprops(labels)
#     ColsLines = [get_line(line.coords) for line in regions if line.bbox[2] - line.bbox[0] > col]
#
#     # 将所有找出的线在黑色图片上画出，找出只有这些线存在的图片上的连通区域
#     tmp = np.zeros(input_size[::-1], dtype='uint8')
#     tmp = draw_lines(tmp, ColsLines + RowsLines, color=255, line_width=1)
#     cv2.imwrite('temp_path/aa.png', tmp)
#
#     labels = measure.label(tmp > 0, connectivity=2)
#     regions = measure.regionprops(labels)
#
#     # 去除这些连通区域中有一边较小的区域以及不在整个大连通区域的线段
#     for region in regions:
#         ymin, xmin, ymax, xmax = region.bbox
#         label = region.label
#         if ymax - ymin < 20 or xmax - xmin < 20:
#             labels[labels == label] = 0
#     labels = measure.label(labels > 0, connectivity=2)
#
#     indY, indX = np.where(labels > 0)
#     xmin, xmax = indX.min(), indX.max()
#     ymin, ymax = indY.min(), indY.max()
#     RowsLines = [p for p in RowsLines if
#                  xmin <= p[0] <= xmax and xmin <= p[2] <= xmax and ymin <= p[1] <= ymax and ymin <= p[3] <= ymax]
#     ColsLines = [p for p in ColsLines if
#                  xmin <= p[0] <= xmax and xmin <= p[2] <= xmax and ymin <= p[1] <= ymax and ymin <= p[3] <= ymax]
#
#     # 将坐标转换成原图中的坐标
#     RowsLines = [[box[0] / fx, box[1] / fy, box[2] / fx, box[3] / fy] for box in RowsLines]
#     ColsLines = [[box[0] / fx, box[1] / fy, box[2] / fx, box[3] / fy] for box in ColsLines]
#     return RowsLines, ColsLines


def find_unclear_lines( row_lines, alph=15 ):
    # 找出不清楚的线并且合并线段
    nrow = len(row_lines)

    newRowsLines = []

    row_used = []

    # 四个点在一条直线上且最近的两个点距离小于 alph
    if nrow > 1:
        for i in range(nrow - 1):
            x1, y1, x2, y2 = row_lines[i]
            if x1 != x2:
                k12 = (y2 - y1) / (x2 - x1)
                a12 = math.degrees(math.atan(k12))
            else:
                a12 = 90

            for j in range(i + 1, nrow):
                x3, y3, x4, y4 = row_lines[j]
                if x3 != x4:
                    k34 = (y4 - y3) / (x4 - x3)
                    a34 = math.degrees(math.atan(k34))
                else:
                    a34 = 90

                # print("a12:{}, a34:{}".format(a12, a34))
                # 如果两条线的角度偏差不超过2°则认为是平行的
                if abs(a34 - a12) < 2:
                    # 计算 a13
                    if x1 != x3:
                        k13 = (y3 - y1) / (x3 - x1)
                        a13 = math.degrees(math.atan(k13))
                    else:
                        a13 = 90

                    # 如果 a12 和 a13相差不超过1°则认为 12和13重合
                    # 求四个点之间的距离，最短距离小于alph则连接，用最长的那条取代
                    if abs(a13 - a12) < 1:
                        dis13 = dist((x1, y1), (x3, y3))
                        dis14 = dist((x1, y1), (x4, y4))
                        dis23 = dist((x2, y2), (x3, y3))
                        dis24 = dist((x2, y2), (x4, y4))
                        dist_list = [dis13, dis14, dis23, dis24]
                        min_dis = min(dist_list)
                        # print("min dis is:", min_dis)

                        if min_dis < alph:
                            max_index = dist_list.index(max(dist_list))
                            # print("max index:", max_index)
                            if max_index == 0:
                                # 返回13
                                newRowsLines.append([x1, y1, x3, y3])
                            elif max_index == 1:
                                # 返回14
                                newRowsLines.append([x1, y1, x4, y4])
                            elif max_index == 2:
                                # 返回23
                                newRowsLines.append([x2, y2, x3, y3])
                            else:
                                # 返回24
                                newRowsLines.append([x2, y2, x4, y4])

                            # 记录已经使用的线段
                            # print("use:", (i, j))
                            row_used.append(i)
                            row_used.append(j)

    for i in range(nrow):
        if i not in row_used:
            newRowsLines.append(row_lines[i])

    return newRowsLines


def adjust_lines( RowsLines, ColsLines, alph=50 ):
    ##调整line

    nrow = len(RowsLines)
    ncol = len(ColsLines)
    newRowsLines = []
    newColsLines = []
    for i in range(nrow):

        x1, y1, x2, y2 = RowsLines[i]
        cx1, cy1 = (x1 + x2) / 2, (y1 + y2) / 2
        for j in range(nrow):
            if i != j:
                x3, y3, x4, y4 = RowsLines[j]
                cx2, cy2 = (x3 + x4) / 2, (y3 + y4) / 2
                if (x3 < cx1 < x4 or y3 < cy1 < y4) or (x1 < cx2 < x2 or y1 < cy2 < y2):
                    continue
                else:
                    r = sqrt((x1, y1), (x3, y3))
                    if r < alph:
                        newRowsLines.append([x1, y1, x3, y3])
                    r = sqrt((x1, y1), (x4, y4))
                    if r < alph:
                        newRowsLines.append([x1, y1, x4, y4])

                    r = sqrt((x2, y2), (x3, y3))
                    if r < alph:
                        newRowsLines.append([x2, y2, x3, y3])
                    r = sqrt((x2, y2), (x4, y4))
                    if r < alph:
                        newRowsLines.append([x2, y2, x4, y4])

    for i in range(ncol):
        x1, y1, x2, y2 = ColsLines[i]
        cx1, cy1 = (x1 + x2) / 2, (y1 + y2) / 2
        for j in range(ncol):
            if i != j:
                x3, y3, x4, y4 = ColsLines[j]
                cx2, cy2 = (x3 + x4) / 2, (y3 + y4) / 2
                if (x3 < cx1 < x4 or y3 < cy1 < y4) or (x1 < cx2 < x2 or y1 < cy2 < y2):
                    continue
                else:
                    r = sqrt((x1, y1), (x3, y3))
                    if r < alph:
                        newColsLines.append([x1, y1, x3, y3])
                    r = sqrt((x1, y1), (x4, y4))
                    if r < alph:
                        newColsLines.append([x1, y1, x4, y4])

                    r = sqrt((x2, y2), (x3, y3))
                    if r < alph:
                        newColsLines.append([x2, y2, x3, y3])
                    r = sqrt((x2, y2), (x4, y4))
                    if r < alph:
                        newColsLines.append([x2, y2, x4, y4])

    return newRowsLines, newColsLines


def get_cell( filepath, tableNet, prob, alph, row, col ):
    """
    从图片中获取所有的单元格
    """
    img = cv2.imread(filepath)
    h, w = img.shape[:2]
    dir_name = os.path.dirname(filepath)

    # 将其转为正方形再传入模型中进行识别
    # width = max(h, w)
    # width = (width // 128 + 1) * 128
    # input_size = (width, width)
    input_size = (1024, 1024)

    RowsLines, ColsLines = get_table_rowcols(img, tableNet, input_size, prob,alph, row, col)
    # # 找出可能不太清楚断掉的直线并且合并线段
    # RowsLines = find_unclear_lines(RowsLines, alph=alph)
    # ColsLines = find_unclear_lines(ColsLines, alph=alph)
    #
    # nrow = len(RowsLines)
    # ncol = len(ColsLines)
    #
    # # 横线竖线相交处不清楚则画出
    # min_dis = 10
    # for i in range(nrow):
    #     for j in range(ncol):
    #         RowsLines[i] = line_line(RowsLines[i], ColsLines[j], min_dis)
    #         ColsLines[j] = line_line(ColsLines[j], RowsLines[i], min_dis)

    # 将找到的所有线都在全黑的图上画出，再在这个图上寻找连通区域即每一个格子
    tmp = np.zeros((h, w), dtype='uint8')
    tmp = draw_lines(tmp, ColsLines + RowsLines, color=255, line_width=1)
    tmp_filepath = '{}/tmp.png'.format(dir_name)
    cv2.imwrite(tmp_filepath, tmp)

    # 找出每一个格子的区域
    tabelLabels = measure.label(tmp == 0, connectivity=2)
    regions = measure.regionprops(tabelLabels)

    rboxes = []
    for region in regions:
        # print("area:", region.bbox_area)
        if region.bbox_area < h * w / 2:
            rbox = minAreaRectBox(region.coords)
            rboxes.append(rbox)

    return rboxes, RowsLines
